
from .dwpose import DWposeDetector
from .util import resize_image,HWC3
import cv2
from PIL import Image
dwprocessor = DWposeDetector()

def get_pose(image_path:str,body_threshold=0.3,other_threshold=0.3,is_discard:bool=True,only_one=True)->Image.Image | None:
    """
    获取Pose，不用手动调整大小
    Args:
     image_path: 图片路径
     body_threshold: 人体置信度阈值
     other_threshold: 其他置信度阈值
     is_discard: 是否丢弃置信度低于阈值的框
     only_one: 是否只返回第一个人的框
    Returns:
     Image.Image | None
    """
    input_image = cv2.imread(image_path)
    input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
    #print(input_image.shape)
    input_image = HWC3(input_image)
    #img = resize_image(input_image, 384)
    img=input_image
    #print(img.shape)
    #cv2.imwrite("../images/222-0.jpg",img)
    detected_map = dwprocessor(img,body_threshold=body_threshold,other_threshold=other_threshold,is_discard=is_discard,only_one=only_one)
    if detected_map is None:
        return None
    pil_pose = Image.fromarray(detected_map)
    #input_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
    #cv2.imwrite("pose.png",detected_map)
    return pil_pose

